package site.teamo.biu.net.client.core.http_https_proxy.mode;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import io.netty.channel.Channel;
import io.netty.handler.codec.http.HttpRequest;
import site.teamo.biu.net.client.core.http_https_proxy.ActionClassifier;
import site.teamo.biu.net.client.core.http_https_proxy.ConnectionHandler;
import site.teamo.biu.net.client.core.http_https_proxy.client.DirectTransferClient;
import site.teamo.biu.net.client.core.http_https_proxy.client.ProxyTransferClient;
import site.teamo.biu.net.client.mapper.ConnectionRuleMapper;
import site.teamo.biu.net.common.constant.ConnectionType;
import site.teamo.biu.net.common.http_https_proxy.ChannelGroups;
import site.teamo.biu.net.common.http_https_proxy.ConnectionState;
import site.teamo.biu.net.client.core.http_https_proxy.TrafficStatistics;
import site.teamo.biu.net.common.http_https_proxy.constant.ActionEnum;
import site.teamo.biu.net.common.util.HostPort;
import site.teamo.biu.net.common.util.HttpUtils;
import site.teamo.biu.tool.spring.util.SpringApplicationContextHolder;


/**
 * @author haocongshun
 * @date 2023/06/13 14:51:33
 */
public class RuleConnectionHandler implements ConnectionHandler {

    private final DirectTransferClient directTransferClient;

    private final ProxyTransferClient proxyTransferClient;

    private final ActionClassifier actionClassifier;

    public RuleConnectionHandler(TrafficStatistics trafficStatistics) {
        this.directTransferClient = new DirectTransferClient(trafficStatistics.getDirectIn(), trafficStatistics.getDirectOut());
        this.proxyTransferClient = new ProxyTransferClient(trafficStatistics.getProxyIn(), trafficStatistics.getProxyOut());
        this.actionClassifier = buildActionClassifier();
    }

    @Override
    public void http(HostPort hostPort, Channel channelFromBrowser, HttpRequest request) {
        ActionEnum action = this.actionClassifier.classify(hostPort);
        switch (action) {
            case REJECT -> channelFromBrowser.writeAndFlush(HttpUtils.forbiddenResponse());
            case DIRECT -> directTransferClient.http(hostPort, channelFromBrowser, request);
            case PROXY -> proxyTransferClient.http(hostPort, channelFromBrowser, request);
        }
    }

    @Override
    public void https(HostPort hostPort, Channel channelFromBrowser) {
        ActionEnum action = this.actionClassifier.classify(hostPort);
        switch (action) {
            case REJECT -> channelFromBrowser.writeAndFlush(HttpUtils.forbiddenResponse());
            case DIRECT -> directTransferClient.https(hostPort, channelFromBrowser);
            case PROXY -> proxyTransferClient.https(hostPort, channelFromBrowser);
        }
    }

    @Override
    public void close() {
        this.directTransferClient.close();
        this.proxyTransferClient.close();
    }

    @Override
    public ConnectionState directConnectionState() {
        ChannelGroups channelGroups = directTransferClient.getChannelGroups();
        ConnectionState connectionState = new ConnectionState();
        channelGroups.forEach((key, value) -> connectionState.add(key, ConnectionType.DIRECT, value.size(), value.getTimestamp()));
        return connectionState;
    }

    @Override
    public ConnectionState proxyConnectionState() {
        ChannelGroups channelGroups = proxyTransferClient.getChannelGroups();
        ConnectionState connectionState = new ConnectionState();
        channelGroups.forEach((key, value) -> connectionState.add(key, ConnectionType.DIRECT, value.size(), value.getTimestamp()));
        return connectionState;
    }

    private ActionClassifier buildActionClassifier() {
        ConnectionRuleMapper connectionRuleMapper = SpringApplicationContextHolder.getBean(ConnectionRuleMapper.class);
        return new ActionClassifier(connectionRuleMapper.selectList(new QueryWrapper<>()), ActionEnum.DIRECT);
    }
}
